import logging
import numpy as np
import dgl
from typing import Dict, List, Set, Tuple, Optional, Any
from source.data_processor.utils.experiment_data import ExperimentData

logger = logging.getLogger(__name__)


class GraphProcessor:
    """
    系统依赖图处理类，负责构建和处理系统依赖图
    """

    def __init__(self):
        """初始化图处理器"""
        pass

    def extract_useful_entities(self, experiment_data: ExperimentData) -> Set[str]:
        """
        从ExperimentData中提取所有有用的实体（服务和容器）

        提取步骤:
        1. 使用ExperimentData的方法获取所有服务
        2. 使用ExperimentData的方法获取所有容器
        3. 合并这些实体ID，确保在异常前后使用相同的实体集合

        参数:
            experiment_data: ExperimentData对象

        返回:
            set: 有用实体ID的集合
        """
        # 从实验数据中获取所有服务和容器
        all_services = experiment_data.get_all_services()
        all_containers = experiment_data.get_all_containers()

        # 合并所有实体
        all_entities = all_services.union(all_containers)

        # 检查是否有实体同时作为服务和容器
        service_container_overlap = all_services.intersection(all_containers)
        if service_container_overlap:
            logger.warning(
                f"发现 {len(service_container_overlap)} 个实体同时作为服务和容器"
            )

        logger.info(
            f"从实验数据中提取了 {len(all_services)} 个服务和 {len(all_containers)} 个容器，共 {len(all_entities)} 个实体"
        )

        return all_entities

    def prepare_graph(
        self,
        interface_deps: Dict[str, List[str]],
        graph_builder,
        useful_entities: Set[str],
    ) -> Tuple:
        """
        准备系统依赖图和节点映射

        处理步骤:
        1. 创建节点列表并建立索引映射
        2. 添加服务间调用关系的边
        3. 添加容器-服务提供关系的边
        4. 创建DGL图并添加自环

        参数:
            interface_deps: 接口依赖关系
            graph_builder: 图构建器
            useful_entities: 有用实体集合

        返回:
            tuple: (图对象, 节点ID列表, ID到索引映射, 索引到ID映射)
        """
        # 步骤1: 创建节点列表并建立索引映射
        # 使用所有有用实体作为节点，并按照字母顺序排序，确保索引顺序稳定
        all_cmdb_ids = sorted(list(useful_entities), key=lambda x: str(x))

        # 创建节点到索引的映射 - 索引顺序与排序后的cmdb_ids一致
        cmdb_id_to_idx = {node: i for i, node in enumerate(all_cmdb_ids)}
        idx_to_cmdb_id = {i: node for i, node in enumerate(all_cmdb_ids)}

        # 初始化边列表
        src_nodes = []
        dst_nodes = []

        # 步骤2: 添加服务间调用关系的边
        service_call_edges = 0
        for caller, callees in interface_deps.items():
            caller_idx = cmdb_id_to_idx.get(caller)
            if caller_idx is None:
                continue

            for callee in callees:
                callee_idx = cmdb_id_to_idx.get(callee)
                if callee_idx is None:
                    continue

                src_nodes.append(caller_idx)
                dst_nodes.append(callee_idx)
                service_call_edges += 1

        # 步骤3: 添加容器-服务提供关系的边
        container_service_edges = 0
        for container, services in graph_builder.container_interfaces.items():
            if container not in useful_entities:
                continue

            container_idx = cmdb_id_to_idx.get(container)
            if container_idx is None:
                continue

            for service in services:
                if service not in useful_entities:
                    continue

                service_idx = cmdb_id_to_idx.get(service)
                if service_idx is None:
                    continue

                # 添加双向边
                src_nodes.append(container_idx)
                dst_nodes.append(service_idx)
                src_nodes.append(service_idx)
                dst_nodes.append(container_idx)
                container_service_edges += 2

        # 步骤4: 创建DGL图并添加自环
        g = dgl.graph((src_nodes, dst_nodes), num_nodes=len(all_cmdb_ids))
        g = dgl.add_self_loop(g)

        logger.info(f"构建图完成，{len(all_cmdb_ids)}个节点，{g.number_of_edges()}条边")
        logger.info(f"节点按照cmdb_id字母顺序排序，确保模型输入一致性")
        return g, all_cmdb_ids, cmdb_id_to_idx, idx_to_cmdb_id
